package token_bucket

import (
	"ratelimit"
	"sync/atomic"
	"time"
)

type TokenBucketLimit struct {
	tokens  chan struct{}
	closeCh chan struct{}
	closed  uint32
}

func New(capacity int, interval time.Duration) ratelimit.Limiter {
	limiter := &TokenBucketLimit{tokens: make(chan struct{}, capacity), closeCh: make(chan struct{}, 1)}
	limiter.init(capacity, interval)
	return limiter
}

func (t *TokenBucketLimit) run(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			select {
			case t.tokens <- struct{}{}:
			case <-t.closeCh: // 如果 tokens 满了，检查 closeCh 是否关闭
				return
			default:

			}
		case <-t.closeCh:
			return
		}
	}
}

func (t *TokenBucketLimit) init(capacity int, interval time.Duration) {
	for i := 0; i < capacity; i++ {
		t.tokens <- struct{}{}
	}
	go t.run(interval)
}

func (t *TokenBucketLimit) Allow() bool {
	select {
	case <-t.tokens:
		return true
	case <-t.closeCh:
		return true
	default:
		return false
	}
}

func (t *TokenBucketLimit) close() {
	if !atomic.CompareAndSwapUint32(&t.closed, 0, 1) {
		return
	}
	close(t.closeCh)
}
